#ifndef VISUALIZE_BBOX_UTIL_H
#define VISUALIZE_BBOX_UTIL_H

#include <glog/logging.h>

#include <fstream>
#ifdef ENABLE_FREETYPE
#include <opencv2/freetype.hpp>
#endif
#include <opencv2/opencv.hpp>
#include <string>
#include <vector>

static void split(const std::string &s, std::vector<std::string> &tokens,
                  const std::string &delimiters = ",") {
  std::string::size_type lastPos = s.find_first_not_of(delimiters, 0);
  std::string::size_type pos = s.find_first_of(delimiters, lastPos);
  while (std::string::npos != pos || std::string::npos != lastPos) {
    tokens.push_back(s.substr(lastPos, pos - lastPos));
    lastPos = s.find_first_not_of(delimiters, pos);
    pos = s.find_first_of(delimiters, lastPos);
  }
}

void parseClassNames(std::vector<std::string> &class_names,
                     const std::string &label_file_path) {
  std::string line;
  std::ifstream label_file(label_file_path);
  if (!label_file.good()) {
    LOG(FATAL) << "Object label file " << label_file_path << " doesn't exist";
  }
  while (!label_file.eof()) {
    std::getline(label_file, line, '\n');
    if (line.length() == 0) continue;
    class_names.push_back(line);
  }
}

void parseClassNamesAndColors(std::vector<std::string> &class_names,
                              std::map<int, cv::Scalar> &class_ids_to_colors,
                              const std::string &label_file_path) {
  std::string line;
  std::ifstream label_file(label_file_path);
  std::vector<std::string> tokens;
  std::string delimiter = ",";
  if (!label_file.good()) {
    LOG(FATAL) << "Object label file " << label_file_path << " doesn't exist";
  }
  int cur_class_id = 0;
  while (!label_file.eof()) {
    std::getline(label_file, line, '\n');
    if (line.length() == 0) continue;
    split(line, tokens, delimiter);
    assert(tokens.size() == 4);  // format class_id,b,g,r
    class_names.push_back(tokens[0]);
    class_ids_to_colors[cur_class_id++] = cv::Scalar(
        std::stoi(tokens[1]), std::stoi(tokens[2]), std::stoi(tokens[3]));
    tokens.clear();
  }
}

#ifdef ENABLE_FREETYPE
class VisualizeBboxUtil {
 public:
  VisualizeBboxUtil() {}
  VisualizeBboxUtil(std::string font_path) { init(font_path); }
  void init(std::string font_path) {
    ft2_ = cv::freetype::createFreeType2();
    ft2_->loadFontData(font_path, 0);
    line_thickness_ = 2;
    font_height_ = 20;
    font_color_ = cv::Scalar::all(255);
    DLOG(INFO) << "Init VisualizeBboxUtil Successfully";
  }

  void drawOneBBox(cv::Mat &frame, float *bbox, std::string &text,
                   cv::Scalar bbox_color) {
    DLOG(INFO) << "Start draw one bbox";
    if (!line_thickness_) {
      line_thickness_ = int(0.002 * (frame.cols + frame.rows) / 2 + 1);
    }

    // draw bbox
    // params from yolov5 plot_one_box
    cv::Rect rect = cv::Rect(int(bbox[0]), int(bbox[1]), int(bbox[2] - bbox[0]),
                             int(bbox[3] - bbox[1]));
    cv::rectangle(frame, rect, bbox_color, line_thickness_, cv::LINE_AA);
    DLOG(INFO) << "Successfully draw bbox";

    // draw chinese text
    // params from yolov5 plot_one_box
    int font_thickness = std::max(1, line_thickness_ - 1);
    ft2_->putText(frame, text, cv::Point(rect.x, rect.y - 2), font_height_,
                  font_color_, font_thickness, cv::LINE_AA, true);
    DLOG(INFO) << "Successfully draw text";
  }

 private:
  cv::Ptr<cv::freetype::FreeType2> ft2_;
  int line_thickness_, font_height_;
  cv::Scalar font_color_;
};
#else
class VisualizeBboxUtil {
 public:
  VisualizeBboxUtil() {}
  VisualizeBboxUtil(std::string font_path) { init(font_path); }
  void init(std::string font_path) {
    line_thickness_ = 2;
    font_face_ = cv::FONT_HERSHEY_SIMPLEX;
    font_scale_ = 1.0;
    // font_color_ = cv::Scalar::all(255);
    font_color_ = cv::Scalar(0, 0, 255);
    DLOG(INFO) << "Init VisualizeBboxUtil Successfully";
  }

  void drawOneBBox(cv::Mat &frame, float *bbox, std::string &text,
                   cv::Scalar bbox_color) {
    DLOG(INFO) << "Start draw one bbox";
    if (!line_thickness_) {
      line_thickness_ = int(0.002 * (frame.cols + frame.rows) / 2 + 1);
    }

    // draw bbox
    // params from yolov5 plot_one_box
    cv::Rect rect = cv::Rect(int(bbox[0]), int(bbox[1]), int(bbox[2] - bbox[0]),
                             int(bbox[3] - bbox[1]));
    cv::rectangle(frame, rect, bbox_color, line_thickness_, cv::LINE_AA);
    DLOG(INFO) << "Successfully draw bbox";

    // draw chinese text
    // params from yolov5 plot_one_box
    int font_thickness = std::max(1, line_thickness_ - 1);
    cv::putText(frame, text, cv::Point(rect.x, rect.y + 30), font_face_,
                font_scale_, font_color_, font_thickness, cv::LINE_AA);
    DLOG(INFO) << "Successfully draw text";
  }

 private:
  int line_thickness_, font_face_;
  double font_scale_;
  cv::Scalar font_color_;
};
#endif

#endif  // VISUALIZE_BBOX_UTIL_H
